import torch
from torchvision import transforms
from torch import optim
import matplotlib.pyplot as plt
from loader import load_CIFAR10
from models import ResNet


def show_pic(acc,loss,config,save=False):
    x1 = range(0, config["epoch"])
    x2 = range(0, config["epoch"])
    y1 = acc
    y2 = loss
    plt.subplot(2, 1, 1)
    plt.plot(x1, y1, 'o-')
    plt.title('Train accuracy')
    plt.ylabel('Train accuracy')
    plt.subplot(2, 1, 2)
    plt.plot(x2, y2, '.-')
    plt.xlabel('Train loss')
    plt.ylabel('Train loss')
    # plt.show()
    if save == True:
        plt.savefig("regularization/loss_lr_{}_batchsize_{}_epoch{}.jpg".format(config["lr"],config["batchsize"],config["epoch"]))

if __name__ == '__main__':

    transform = transforms.ToTensor()
    train_data = load_CIFAR10("../datasets/",train = True, transforms=transform) #data and label as tensor

    config = {"epoch": 20,
            "batchsize": 32,
            "lr": 2e-4 #2e-4这个学习率挺不错
            }
    print("lr:{}".format(config["lr"]))
    net = ResNet.ResNet(ResNet.Residual_Block,[2,2,2,2],10)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    opt = optim.Adam(params=net.parameters(), lr=config["lr"],weight_decay=0.003)#0.003 good 
    scheduler = torch.optim.lr_scheduler.StepLR(opt,step_size=5,gamma=0.3)#0.3 good

    Accuracy_list, Loss_list = net.train(train_data,config,opt,device)

    test_data = load_CIFAR10("../datasets/",train = False, transforms=transform) #data and label as tensor
    acc_list = []

    for i in range(20):
        net = ResNet.ResNet(ResNet.Residual_Block,[2,2,2,2],10)
        net.load_state_dict(torch.load('parameters/regularization/model_parameter_lr0.0002_epoch{}.pt'.format(i)))   
        
        acc = net.run(test_data,device)
        acc_list.append(acc)
    
    x1 = range(0, config["epoch"])
    x2 = range(0, config["epoch"])
    y1 = Accuracy_list
    y2 = Loss_list
    y3 = acc_list
    plt.subplot(2, 1, 1)
    plt.plot(x1, y1, 'o-')
    plt.plot(x1, y3, 'o-',color="green")
    plt.title('Train accuracy')
    plt.ylabel('Train accuracy')
    plt.subplot(2, 1, 2)
    plt.plot(x2, y2, '.-')
    plt.xlabel('epoch')
    plt.ylabel('Train loss')

    plt.savefig("regularization/acc_loss_lr_{}_batchsize_{}_lambda0.003.jpg".format(config["lr"],config["batchsize"]))

